#!/usr/bin/python
# -*- coding: utf-8 -*-
# @File : Helper_rabbitmq
# @Time : 2023/10/8 15:40
# @Author : yu.qian
"""
 ____    _____    ____   _                       _   ____     ___    ____    ____  
|  _ \  |_   _|  / ___| | |   ___    _   _    __| | |___ \   / _ \  |___ \  | ___| 
| | | |   | |   | |     | |  / _ \  | | | |  / _` |   __) | | | | |   __) | |___ \ 
| |_| |   | |   | |___  | | | (_) | | |_| | | (_| |  / __/  | |_| |  / __/   ___) |
|____/    |_|    \____| |_|  \___/   \__,_|  \__,_| |_____|  \___/  |_____| |____/ 
                                                                                   
"""
import uuid
import random
import threading
import logging
import pika

_logger = logging.getLogger(__name__)

# TODO: 请在此处配置RabbitMQ节点（实际使用时可以写在Config中）,当前只是实现一个rabbitmq的上下文管理器
RABBITMQ_NODES = []

class RabbitmqHelper(object):
    def __init__(self, name=None):
        if name:
            self.name = name
        else:
            self.name = str(uuid.uuid4())
        self.rabbitmq_pool = RabbitMQPool()

    def __enter__(self):
        return self.rabbitmq_pool.get_connection(self.name)

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.rabbitmq_pool.return_connection(self.name)

class RabbitMQPool:
    __instance = None
    __lock = threading.Lock()

    def __init__(self):
        self.nodes = [pika.URLParameters(node) for node in RABBITMQ_NODES]
        self.connect_pool = {}

    def __new__(cls):
        if not cls.__instance:
            cls.__instance = object.__new__(cls)
        return cls.__instance

    def create_connection(self, name):
        try:
            self.__lock.acquire()
            if name not in self.connect_pool:
                random.shuffle(self.nodes)
                self.connect_pool[name] = pika.BlockingConnection(self.nodes)
                _logger.debug(f"create connection name: {name}")
        finally:
            self.__lock.release()
        return self.connect_pool[name]

    def remove_connection(self, name):
        try:
            self.__lock.acquire()
            if name in self.connect_pool:
                connection = self.connect_pool[name]
                del self.connect_pool[name]
                connection.close()
        finally:
            self.__lock.release()

    def get_connection(self, name=None):
        if name in self.connect_pool:
            connection = self.connect_pool[name]
        else:
            connection = self.create_connection(name)

        if connection.is_open:
            return connection
        else:
            self.remove_connection(name)
            return self.create_connection(name)

    def return_connection(self, name=None):
        if self.connect_pool[name].is_closed:
            self.remove_connection(name)

if __name__ == '__main__':
    # 下面是具体使用例子
    import json
    from packages.distask import app

    @app.task(ignore_result=True)
    def task_send_ext_params(ext_params:dict):

        with RabbitmqHelper() as connection:
            # 建立连接并创建频道:
            channel = connection.channel()
            # 创建一个交换机
            # 声明exchange，由exchange指定消息在哪个队列传递，如不存在，则创建.durable = True 代表exchange持久化存储，False 非持久化存储
            exchange_name = 'zyf.survey'
            channel.exchange_declare(exchange=exchange_name, exchange_type='direct', durable=True)

            # 创建一个队列
            # 声明消息队列，消息将在这个队列传递，如不存在，则创建。durable = True 代表消息队列持久化存储，False 非持久化存储
            queue_name = 'zyf_survey'
            channel.queue_declare(queue_name, durable=True)

            # 将队列绑定到交换机:exchange --direct--> queue
            channel.queue_bind(exchange=exchange_name, queue=queue_name, routing_key=queue_name)
            # 声明消息持久化，从内存转存到硬盘 delivery_mode = 2 声明消息在队列中持久化，delivery_mod = 1 消息非持久化
            properties = pika.BasicProperties(delivery_mode=2)
            # 发送消息
            channel.basic_publish(
                exchange=exchange_name,
                routing_key=queue_name,
                properties=properties,
                body=json.dumps(ext_params),
            )
